#!/usr/bin/env python


import roslib; roslib.load_manifest('knobs_rev')
import rospy
import message_filters
import image_geometry
import actionlib

from smach import State, StateMachine
from sensor_msgs.msg import Image, CameraInfo, PointCloud2
from arm_navigation_msgs.msg import MoveArmGoal, MoveArmAction, JointConstraint
from pr2_controllers_msgs.msg import Pr2GripperCommandGoal, Pr2GripperCommandAction

from std_msgs.msg import Float32

from knobs_rev.msg import PoseArmAction, PoseArmGoal

import cv
from cv_bridge import CvBridge, CvBridgeError

from knobs_rev.srv import PlaneCloud, PlaneCloudRequest, PlaneCloudResponse

from PySide.QtCore import * 
from PySide.QtGui import * 

import sys
from time import sleep


class ClearArm(State):
    def __init__(self):
        State.__init__(self, outcomes=['succeeded', 'failed'])
        
        self.right_arm_client = actionlib.SimpleActionClient('pose_arm', PoseArmAction)
        self.right_arm_client.wait_for_server()

    def execute(self, userdata):
        goal = PoseArmGoal()
        goal.pose.header.frame_id = 'torso_lift_link'
        goal.pose.pose.position.x = 0.5
        goal.pose.pose.position.y = -0.3
        goal.pose.pose.position.z = -0.1

        self.right_arm_client.send_goal(goal)
        self.right_arm_client.wait_for_result()

        return 'succeeded'


class Click(State):
    def __init__(self, group_name, buttons):
        State.__init__(self, outcomes=['succeeded', 'failed'])
        self.layout = QVBoxLayout()

        self.button_list = [QPushButton(name) for name in buttons]
        self.return_list = [value for value in buttons.values()]
        map(lambda x: x[0].clicked.connect(lambda: self.callback(x[1])), zip(self.button_list, self.return_list))
        map(self.layout.addWidget, self.button_list)

        self.group = QGroupBox(group_name)
        self.group.setLayout(self.layout)

    def callback(self, value):
        self.return_value = value
        self.group.hide()
        app.quit()

    def execute(self, userdata):
        self.return_value = None

        self.group.show()
        app.exec_()

        return self.return_value


class Display(State):
   def __init__(self, feed):
      # Smach
      smach.State.__init__(self, outcomes=['success', 'aborted'],
                       output_keys=['low', 'high'])
      # Flags and global vars
      self.qimage = QImage()
      # OpenCv setup
      self.bridge = CvBridge()
      # Publishers and Subscribers
      self.image_sub = rospy.Subscriber(feed, Image, self.callback)
      rospy.loginfo("Waiting for transform from base_link to wide_stereo_optical_frame")
      # Pyside
      self.imageReady = QSignal()
      self.imageLock = threading.Lock()
      self.group = QGroupBox('group1')
      self.label2 = QLabel()

      self.button = QPushButton('done')

      vbox = QVBoxLayout()
      vbox.addWidget(self.label2)
      vbox.addWidget(self.button)
      self.group.setLayout(vbox)

   def done(self):
      self.group.hide()
      app.quit()
   def execute(self, userdata):
      self.imageReady.signal.connect(self.showImage)
      self.button.clicked.connect(self.done)
      self.group.show()
      app.exec_()
      rospy.loginfo("app executed")
      return 'success' 

   def showImage(self):
      rospy.loginfo("showImage")

      self.imageLock.acquire()
      qpixmap = QPixmap.fromImage(self.qimage)
      self.label2.setPixmap(qpixmap)
      self.imageLock.release()

      
   def callback(self, data):
      rospy.loginfo("Received image")
      if rospy.is_shutdown():
         return
      try:
         cv_image = self.bridge.imgmsg_to_cv(data, "rgba8")
      except CvBridgeError, e:
         print e
         return
      self.imageLock.acquire()
      self.qimage = IplQImage(cv_image)
      self.imageLock.release()
      self.imageReady.signal.emit()



class CalculatePlane(State):
    def __init__(self, cloud_name, camera_info_name):
        State.__init__(self,
                       outcomes=['succeeded', 'failed'],
                       input_keys=['low', 'high'],
                       output_keys=['normal', 'point'])
        self.cloud_name = cloud_name
        self.camera_info_name = camera_info_name
        self.done = False

    def cloud_callback(self, cloud, camera_info, low, high):
        if self.done:
            return

        # Fill the service request
        request = PlaneCloudRequest()
        request.cloud = cloud

        # Convert from image space to world space for winnowing out the point cloud.
        camera_model = image_geometry.PinholeCameraModel()
        camera_model.fromCameraInfo(camera_info)
        ul_ray = camera_model.projectPixelTo3dRay((low[0], low[1]))
        ur_ray = camera_model.projectPixelTo3dRay((high[0], low[1]))
        ll_ray = camera_model.projectPixelTo3dRay((low[0], high[1]))
        lr_ray = camera_model.projectPixelTo3dRay((high[0], high[1]))
                                                  
        # Fake the min and max z values for now
        request.low_z = 0.5
        request.high_z = 1.5

        # Calculate the the x and y values, based on the closest z value
        request.low_x = min(ul_ray[0], ll_ray[0]) * request.low_z
        request.high_x = max(ur_ray[0], lr_ray[0]) * request.low_z
        request.low_y = min(ul_ray[1], ur_ray[1]) * request.low_z
        request.high_y = min(ll_ray[1], lr_ray[1]) * request.low_z

        # Hack!
        request.low_x = -100.0
        request.high_x = 100.0
        request.low_y = -100.0
        request.high_y = 100.0
        request.low_z = -100.0
        request.high_z = 100.0
        
        # Get a service proxy for the plane segmenter
        plane_seg = rospy.ServiceProxy('plane_finder', PlaneCloud)
        
        # Get the plane from the point cloud
        response = PlaneCloudResponse()
        response = plane_seg(request)

        # Calculate the ray from center of selected area
        ray = camera_model.projectPixelTo3dRay(((low[0] + high[0]) / 2.0, (low[1] + high[1]) / 2.0))

        # Intersection of plane and line (assuming that everyone is in the same frame)
        t = response.d / (response.a * ray[0] + response.b * ray[1] + response.c * ray[2])

        # Point on the end of the knob
        x = ray[0] * t;
        y = ray[1] * t;
        z = ray[2] * t;

        # Store the data to pass on to the next state
        self.normal = (response.a, response.b, response.c)
        self.point = (x, y, z)
        print 'Frame:', 

        self.done = True

    def execute(self, userdata):
        print 'low', userdata.low
        print 'high', userdata.high

        cloud_sub = message_filters.Subscriber(self.cloud_name, PointCloud2)
        camera_info_sub = message_filters.Subscriber(self.camera_info_name, CameraInfo)
        sync = message_filters.TimeSynchronizer([cloud_sub, camera_info_sub], 10)
        sync.registerCallback(lambda x,y: self.cloud_callback(x, y, userdata.low, userdata.high))

        while not self.done:
            sleep(0.01)

        userdata.normal = self.normal
        userdata.point = self.point

        return 'succeeded'


class PositionArm(State):
    def __init__(self):
        State.__init__(self,
                       outcomes=['succeeded', 'failed'],
                       input_keys=['normal', 'point'])

        self.right_arm_client = actionlib.SimpleActionClient('pose_arm', PoseArmAction)
        self.right_arm_client.wait_for_server()

    def execute(self, userdata):
        print 'Normal:', userdata.normal
        print 'Point:', userdata.point

        t = -0.2

        goal = PoseArmGoal()
        goal.pose.header.frame_id = 'torso_lift_link'
        goal.pose.pose.position.x = 0.72
        goal.pose.pose.position.y = -0.18
        goal.pose.pose.position.z = 0.09

#        goal.pose.header.frame_id = '/narrow_stereo_optical_frame'
#        goal.pose.header.stamp = rospy.Time.now()
#        goal.pose.pose.position.x = userdata.point[0] #+ t * userdata.normal[0]
#        goal.pose.pose.position.y = userdata.point[1] #+ t * userdata.normal[1]
#        goal.pose.pose.position.z = userdata.point[2] #+ t * userdata.normal[2]

        self.right_arm_client.send_goal(goal)
        self.right_arm_client.wait_for_result()

        return 'succeeded'


class FinalApproach(State):
    def __init__(self):
        State.__init__(self, outcomes = ['succeeded', 'failed'])

        self.right_arm_client = actionlib.SimpleActionClient('/move_right_arm', MoveArmAction)
        self.right_arm_client.wait_for_server()


    def execute(self, userdata):
        goal = MoveArmGoal()
        goal.motion_plan_request.group_name = 'right_arm'
        goal.motion_plan_request.num_planning_attempts = 1
        goal.motion_plan_request.allowed_planning_time = rospy.Duration(5.0)
        goal.motion_plan_request.planner_id = ''
        goal.planner_service_name = 'ompl_planning/plan_kinematic_path'

        joint_angles = {'r_upper_arm_roll_joint':-1.5466936482606153, 
                        'r_shoulder_pan_joint':-0.24061559789458131, 
                        'r_shoulder_lift_joint':-0.13052543328701074, 
                        'r_forearm_roll_joint':-2.8959025637550835, 
                        'r_elbow_flex_joint':-0.67977712040984106, 
                        'r_wrist_flex_joint':-0.50094827692654909, 
                        'r_wrist_roll_joint':-1.7652659302215628} 

        for joint,angle in joint_angles.iteritems():
            j = JointConstraint()
            j.joint_name = joint
            j.position = angle
            j.tolerance_above = 0.01
            j.tolerance_below = 0.01
            j.weight = 1.0
            goal.motion_plan_request.goal_constraints.joint_constraints.append(j)

        # Sometimes we get it, sometimes we don't
        for i in range(10):
            self.right_arm_client.send_goal(goal)
            self.right_arm_client.wait_for_result()
        
        return 'succeeded'


class Gripper(State):
    def __init__(self, close):
        State.__init__(self, outcomes=['succeeded', 'failed'])
        self.close = close

        self.gripper_client = actionlib.SimpleActionClient('r_gripper_controller/gripper_action', Pr2GripperCommandAction)
        self.gripper_client.wait_for_server()

    def execute(self, userdata):
        goal = Pr2GripperCommandGoal()
        if self.close:
            goal.command.position = 0.03
        else:
            goal.command.position = 0.08

        goal.command.max_effort = 50.0

        self.gripper_client.send_goal(goal)
        self.gripper_client.wait_for_result()

        return 'succeeded'


class TurnKnob(State):
    def __init__(self):
        State.__init__(self, outcomes=['succeeded', 'failed'])

        self.layout = QVBoxLayout()
        self.dial = QDial()
        self.dial.valueChanged.connect(self.callback)
        self.dial.setNotchesVisible(True)
        self.dial.setWrapping(True)
        self.layout.addWidget(self.dial)

        self.button = QPushButton('done')
        self.button.clicked.connect(self.done)
        self.layout.addWidget(self.button)

        self.group = QGroupBox('knob control')
        self.group.setLayout(self.layout)

        self.pub = rospy.Publisher('knob/value', Float32)

    def callback(self, value):
        self.pub.publish(float(value) / 99.0 * 3.14159 * 2.0)

    def done(self):
        self.group.hide()
        app.quit()

    def execute(self, userdata):
        self.group.show()
        app.exec_()

        return 'succeeded'


if __name__ == '__main__':
    rospy.init_node('finder')

    app = QApplication(sys.argv)

    machine = StateMachine(outcomes=['succeeded', 'failed'])
    with machine:
        StateMachine.add('CLEAR_ARM', ClearArm(), transitions={'succeeded':'START'})
        StateMachine.add('CLEAR_ARM', ClearArm(), transitions={'succeeded':'SELECT_REGION'})
        StateMachine.add('START', Click('Find', {'knobs':'succeeded', 'quit':'failed'}), transitions={'succeeded':'SELECT_REGION', 'failed':'failed'})
        StateMachine.add('SELECT_REGION', Display('/wide_stereo/right/image_rect_color'), transitions={'succeeded':'CALCULATE_PLANE', 'abort':'succeeded'})
        StateMachine.add('CALCULATE_PLANE', CalculatePlane('/narrow_stereo/points2', '/narrow_stereo/right/camera_info'), transitions={'succeeded':'APPROACH_KNOB'})
        StateMachine.add('APPROACH_KNOB', PositionArm(), transitions={'succeeded':'FINAL_APPROACH'})
        StateMachine.add('FINAL_APPROACH', FinalApproach(), transitions={'succeeded':'GRIP_KNOB'})
        StateMachine.add('GRIP_KNOB', Gripper(True), transitions={'succeeded':'RELEASE_KNOB'})
        StateMachine.add('TURN_KNOB', TurnKnob(), transitions={'succeeded':'RELEASE_KNOB'})
        StateMachine.add('RELEASE_KNOB', Gripper(False), transitions={'succeeded':'CLEAR_ARM'})


    try:
      rospy.loginfo("Starting state machine execution")
      machine.execute()
      rospy.spin()
    except KeyboardInterrupt:
      rospy.loginfo("Shutting down")
    
